# encoding: utf-8
# @File  : main.py
# @Author: shaoyun
# @Date  :  2025/05/11
import asyncio
import os
import platform
import signal
import socket
import sys
import uuid
from datetime import datetime, timezone

import httpx
import psutil
from loguru import logger

# 添加项目根目录到sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.config import settings
from app.db.redis_manager import RedisClient, AsyncRedisPubSub
from app.db.mongo_manager import MongoDB
from app.services.spider_runner import SpiderRunner

# 配置日志
# 添加控制台输出
logger.add(sys.stdout, level="INFO", format="{time} - {level} - {message}")

# 添加文件输出，文件按日期命名
log_file = f"worker_{datetime.now().strftime('%Y%m%d')}.log"
logger.add(log_file, level="INFO", rotation="00:00", retention="7 days", encoding="utf-8",
           format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}")

# 全局变量
NODE_ID = os.environ.get("NODE_ID") or str(uuid.uuid4())
MASTER_URL = os.environ.get("MASTER_URL") or "http://localhost:8000"
NODE_NAME = os.environ.get("NODE_NAME") or f"{platform.node()}-{NODE_ID[:8]}"
IS_RUNNING = True
CURRENT_TASK = None
AUTH_TOKEN = None  # 存储认证token


async def get_auth_token():
    """获取认证token"""
    global AUTH_TOKEN

    if AUTH_TOKEN:
        return AUTH_TOKEN

    # 使用默认的worker凭据登录
    url = f"{MASTER_URL}/api/v1/auth/token"
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                url,
                data={
                    "username": "admin",  # 使用管理员账号或创建专门的worker账号
                    "password": "admin123",
                    "grant_type": "password"
                },
                timeout=10.0
            )

            if response.status_code == 200:
                result = response.json()
                AUTH_TOKEN = result["data"]["access_token"]
                logger.info("Authentication successful")
                return AUTH_TOKEN
            else:
                logger.error(f"Authentication failed: {response.status_code}")
                return None
    except Exception as e:
        logger.error(f"Error getting auth token: {e}")
        return None


async def send_heartbeat():
    """向Master发送心跳"""
    url = f"{MASTER_URL}/api/v1/nodes/{NODE_ID}/heartbeat"
    timestamp = int(datetime.now(timezone.utc).timestamp())

    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                url,
                json={
                    "node_id": NODE_ID,
                    "timestamp": timestamp
                },
                timeout=5.0
            )

            if response.status_code != 200:
                logger.warning(f"Failed to send heartbeat: {response.status_code} - {response.text}")
                return False

            logger.debug("Heartbeat sent successfully")
            return True
    except Exception as e:
        logger.error(f"Error sending heartbeat: {e}")
        return False


async def send_status():
    """向Master发送节点状态"""
    url = f"{MASTER_URL}/api/v1/nodes/{NODE_ID}/status"

    try:
        # 获取系统信息
        cpu_usage = psutil.cpu_percent(interval=1)
        memory = psutil.virtual_memory()
        memory_usage = memory.percent
        disk = psutil.disk_usage('/')
        disk_usage = {
            "total": disk.total,
            "used": disk.used,
            "free": disk.free,
            "percent": disk.percent
        }

        # Windows系统不支持getloadavg，使用cpu_percent代替
        if platform.system() == 'Windows':
            load_avg = [cpu_usage / 100, cpu_usage / 100, cpu_usage / 100]
        else:
            load_avg = list(psutil.getloadavg())

        status_data = {
            "node_id": NODE_ID,
            "cpu_usage": cpu_usage,
            "memory_usage": memory_usage,
            "disk_usage": disk_usage,
            "load_average": load_avg,
            "timestamp": int(datetime.now(timezone.utc).timestamp())
        }

        async with httpx.AsyncClient() as client:
            response = await client.post(
                url,
                json=status_data,
                timeout=5.0
            )

            if response.status_code != 200:
                logger.warning(f"Failed to send status: {response.status_code} - {response.text}")
                return False

            result = response.json()

            # 检查是否有异常报告
            if result.get("is_anomaly", False):
                logger.warning(f"Anomaly detected: {result.get('message', 'Unknown')}")

            logger.debug("Status sent successfully")
            return True
    except Exception as e:
        logger.error(f"Error sending status: {e}")
        return False


async def process_task(task_id: str):
    """处理任务"""
    global CURRENT_TASK
    token = await get_auth_token()

    if not token:
        logger.error("No auth token available")
        return False

    try:
        CURRENT_TASK = task_id
        logger.info(f"Processing task: {task_id}")

        # 从MongoDB获取任务信息
        task = await MongoDB.find_one("tasks", {"_id": task_id})

        if not task:
            logger.error(f"Task {task_id} not found in MongoDB")
            return False

        # 更新任务状态为运行中
        async with httpx.AsyncClient() as client:
            response = await client.put(
                f"{MASTER_URL}/api/v1/tasks/{task_id}",
                json={"status": "running"},
                headers={"Authorization": f"Bearer {token}"},
                timeout=5.0
            )

            if response.status_code != 200:
                logger.warning(f"Failed to update task status: {response.status_code} - {response.text}")

        # 运行爬虫脚本
        spider_id = task.get("spider_id")
        parameters = task.get("parameters", {})

        success, result = await SpiderRunner.run_spider(
            task_id=task_id,
            spider_id=spider_id,
            parameters=parameters
        )

        # 更新任务状态为完成或失败
        new_status = "completed" if success else "failed"

        async with httpx.AsyncClient() as client:
            response = await client.put(
                f"{MASTER_URL}/api/v1/tasks/{task_id}",
                json={"status": new_status},
                headers={"Authorization": f"Bearer {token}"},
                timeout=5.0
            )

            if response.status_code != 200:
                logger.warning(f"Failed to update task status: {response.status_code} - {response.text}")

        logger.info(f"Task {task_id} completed with status: {new_status}")
        return success

    except Exception as e:
        logger.exception(f"Error processing task {task_id}: {e}")

        # 更新任务状态为失败
        async with httpx.AsyncClient() as client:
            try:
                response = await client.put(
                    f"{MASTER_URL}/api/v1/tasks/{task_id}",
                    json={"status": "failed"},
                    headers={"Authorization": f"Bearer {token}"},
                    timeout=5.0
                )
            except Exception as update_error:
                logger.error(f"Failed to update task status after error: {update_error}")

        return False
    finally:
        CURRENT_TASK = None


async def register_node():
    """注册节点到Master"""
    token = await get_auth_token()

    if not token:
        logger.error("Cannot register node without auth token")
        return False

    url = f"{MASTER_URL}/api/v1/nodes/"

    try:
        # 获取本机IP地址
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
    except Exception:
        ip = "127.0.0.1"

    data = {
        "name": NODE_NAME,
        "ip": ip,
        "port": "8001",  # Worker端口
        "is_master": False,
        "description": f"Worker node on {platform.node()} ({platform.system()} {platform.release()})"
    }

    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                url,
                json=data,
                headers={"Authorization": f"Bearer {token}"},
                timeout=10.0
            )

            if response.status_code == 201:
                result = response.json()
                global NODE_ID
                NODE_ID = result.get("id", NODE_ID)
                logger.info(f"Node registered successfully with ID: {NODE_ID}")
                return True
            else:
                logger.error(f"Failed to register node: {response.status_code} - {response.text}")
                return False
    except Exception as e:
        logger.error(f"Error registering node: {e}")
        return False


async def fetch_task():
    """从Redis队列获取任务"""
    try:
        # 先尝试从优先级队列获取
        task_id = RedisClient.pop_priority_task()
        if task_id:
            logger.info(f"Fetched priority task: {task_id}")
            return task_id

        # 再从普通队列获取
        task_id = RedisClient.pop_task(timeout=1)
        if task_id:
            logger.info(f"Fetched task: {task_id}")
            return task_id

        return None
    except Exception as e:
        logger.error(f"Error fetching task: {e}")
        return None


async def heartbeat_loop():
    """心跳循环"""
    while IS_RUNNING:
        try:
            await send_heartbeat()
            await asyncio.sleep(settings.NODE_HEARTBEAT_INTERVAL)
        except Exception as e:
            logger.error(f"Error in heartbeat loop: {e}")
            await asyncio.sleep(5)  # 出错后等待较短时间再次尝试


async def status_loop():
    """状态上报循环"""
    while IS_RUNNING:
        try:
            await send_status()
            await asyncio.sleep(60)  # 每分钟上报一次状态
        except Exception as e:
            logger.error(f"Error in status loop: {e}")
            await asyncio.sleep(10)  # 出错后等待较短时间再次尝试


async def task_loop():
    """任务处理循环"""
    while IS_RUNNING:
        try:
            if CURRENT_TASK is None:
                # 如果当前没有正在处理的任务，尝试获取新任务
                task_id = await fetch_task()
                if task_id:
                    # 启动一个新的协程来处理任务，不阻塞当前循环
                    asyncio.create_task(process_task(task_id))

            await asyncio.sleep(1)  # 短暂等待再次检查
        except Exception as e:
            logger.error(f"Error in task loop: {e}")
            await asyncio.sleep(5)  # 出错后等待较短时间再次尝试


async def subscribe_loop():
    """订阅Redis事件"""
    pubsub_helper = AsyncRedisPubSub(["file.upload", "spider.sync"])

    try:
        await pubsub_helper.start()

        async for message in pubsub_helper.listen():
            if not IS_RUNNING:
                break

            channel = message['channel']
            data = message['data']

            if channel in ["file.upload", "spider.sync"]:
                spider_id = data.get("spider_id")
                file_id = data.get("file_id")

                # 预下载爬虫文件
                try:
                    await SpiderRunner.download_spider_file(spider_id, file_id)
                    logger.info(f"Pre-downloaded spider {spider_id}")
                except Exception as e:
                    logger.error(f"Failed to download spider {spider_id}: {e}")

    finally:
        await pubsub_helper.stop()
async def main():
    """主函数"""
    global IS_RUNNING

    # 初始化
    try:
        # 连接到MongoDB
        await MongoDB.connect_async()
        logger.info("Connected to MongoDB")

        # 注册节点
        if not await register_node():
            logger.warning("Failed to register node, continuing without registration")
    except Exception as e:
        logger.error(f"Initialization error: {e}")
        return

    # 启动各个循环
    heartbeat_task = asyncio.create_task(heartbeat_loop())
    status_task = asyncio.create_task(status_loop())
    task_processing = asyncio.create_task(task_loop())

    # 添加订阅任务
    subscribe_task = asyncio.create_task(subscribe_loop())

    try:
        # 等待所有任务完成（实际上它们会一直运行）
        await asyncio.gather(
            heartbeat_task,
            status_task,
            task_processing,
            subscribe_task
        )
    except asyncio.CancelledError:
        logger.info("Worker is shutting down...")
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
    finally:
        # 清理
        IS_RUNNING = False

        # 取消所有任务
        for task in [heartbeat_task, status_task, task_processing]:
            if not task.done():
                task.cancel()

        # 关闭数据库连接
        await MongoDB.close_async()
        RedisClient.close()

        logger.info("Worker shutdown complete")


def signal_handler(sig, frame):
    """信号处理函数"""
    global IS_RUNNING
    logger.info(f"Received signal {sig}, shutting down...")
    IS_RUNNING = False


if __name__ == "__main__":
    # 注册信号处理
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # 根据平台配置事件循环
    if platform.system() == 'Windows':
        # Windows系统使用ProactorEventLoop
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

    # 运行主程序
    asyncio.run(main())